"""
测试训练好的模型
"""
import tensorflow as tf
from video_classification.data_manage import read_data
from video_classification import config


MODEL_CLIP = config.MODEL_CLIP
MODEL_WIDTH = config.MODEL_WIDTH
MODEL_HEIGHT = config.MODEL_HEIGHT
MODEL_CHANNEL = config.MODEL_CHANNEL
BATCH_SIZE = config.BATCH_SIZE
MODEL_DIR = config.MODEL_DIR
OUTPUT_LIST = config.OUTPUT_LIST


if __name__ == '__main__':
    # 读取数据集和标签
    test_data, test_label, test_classes = read_data(OUTPUT_LIST[1], (MODEL_CLIP, MODEL_WIDTH, MODEL_HEIGHT))
    with tf.Session() as sess:
        # 导入模型
        saver = tf.train.import_meta_graph(MODEL_DIR + '.meta')
        saver.restore(sess, tf.train.latest_checkpoint('models/'))
        # 导入计算图
        graph = tf.get_default_graph()
        x = graph.get_tensor_by_name("input/x:0")
        # 模型测试
        output = []     # 输出结果列表
        corrected = 0   # 正确结果
        # 计算batch次数
        if len(test_data) % BATCH_SIZE == 0:
            batch_num = int(len(test_data) / BATCH_SIZE)
        else:
            batch_num = int(len(test_data) / BATCH_SIZE) + 1
        for i in range(batch_num):
            if (i + 1) * BATCH_SIZE < len(test_data):
                feed_dic = {x: test_data[i * BATCH_SIZE: (i + 1) * BATCH_SIZE]}
            else:
                feed_dic = {x: test_data[i * BATCH_SIZE: len(test_data)]}
            logits = graph.get_tensor_by_name("logits_value:0")
            classification_result = sess.run(logits, feed_dic)
            output.extend(tf.argmax(classification_result, 1).eval())
        # 输出测试结果
        for i in range(len(test_label)):
            print(test_label[i], output[i])
            if int(test_label[i]) == int(output[i]):
                corrected += 1
        print('Acc: %.4f' % (corrected / len(test_label)))
